package com.socket.secure.util;

import cn.hutool.core.util.HexUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.CryptoException;
import cn.hutool.crypto.Mode;
import cn.hutool.crypto.Padding;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.symmetric.AES;
import com.socket.secure.constant.RequsetTemplate;
import com.socket.secure.constant.SecureConstant;
import com.socket.secure.exception.InvalidRequestException;

import javax.servlet.http.HttpSession;

/**
 * AES decryption tool
 */
public class AESUtil {
    /**
     * Generate hex AES key
     *
     * @return AES key
     */
    public static String generateAesKey() {
        return HexUtil.encodeHexStr(SecureUtil.generateKey("AES").getEncoded());
    }

    /**
     * AES encryption (using the current session's key)
     *
     * @param plaintext plaintext
     * @param session   {@linkplain HttpSession}
     * @return ciphertext
     */
    public static String encrypt(String plaintext, HttpSession session) {
        return encrypt(plaintext, getAesKey(session));
    }

    /**
     * AES encryption (custom key)
     *
     * @param plaintext plaintext
     * @return ciphertext
     */
    public static String encrypt(String plaintext, String key) {
        Assert.notNull(key, RequsetTemplate.AESKEY_NOT_FOUNT, InvalidRequestException::new);
        try {
            return getAES(key).encryptBase64(plaintext.getBytes());
        } catch (CryptoException e) {
            throw new InvalidRequestException(RequsetTemplate.AES_ENCRYPT_ERROR, plaintext);
        }
    }

    /**
     * Get the AES key exchanged in the current session, or null if the key does not exist
     */
    public static String getAesKey(HttpSession session) {
        return (String) session.getAttribute(SecureConstant.AESKEY);
    }

    /**
     * AES decryption (using the current session's key)
     *
     * @param ciphertext &lt;ciphertext&gt;
     * @param session    {@linkplain HttpSession}
     * @return plaintext
     */
    public static String decrypt(String ciphertext, HttpSession session) {
        if (StrUtil.isWrap(ciphertext, "<", ">")) {
            return decrypt(StrUtil.unWrap(ciphertext, "<", ">"), getAesKey(session));
        }
        return ciphertext;
    }

    /**
     * AES decryption (custom key)
     *
     * @param ciphertext ciphertext
     * @return plaintext
     */
    public static String decrypt(String ciphertext, String key) {
        Assert.notNull(key, RequsetTemplate.AESKEY_NOT_FOUNT, InvalidRequestException::new);
        if (StrUtil.isNotEmpty(ciphertext)) {
            try {
                return getAES(key).decryptStr(ciphertext);
            } catch (CryptoException e) {
                throw new InvalidRequestException(RequsetTemplate.AES_DECRYPT_ERROR, ciphertext);
            }
        }
        return StrUtil.EMPTY;
    }

    /**
     * Get the built-in AES object
     */
    public static AES getAES(String key) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < key.length() / 2; i++) {
            sb.append(Integer.toString(key.charAt(i), 16));
        }
        byte[] iv = HexUtil.decodeHex(sb.toString());
        return new AES(Mode.CBC, Padding.PKCS5Padding, StrUtil.bytes(key), iv);
    }
}
